package _11_binary_search_tree;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

/**
 * @author: mornd
 * @dateTime: 2023/6/12 - 16:33
 * 二叉搜索树定义，key 可以为其它可比较类型，这里定义为 int
 */
public class BSTree<E> {
    // 树根
    public BSTNode<E> root;

    /**
     * 查找关键字对应的值
     *
     * @param key
     * @return
     */
    public E get(int key) {
        BSTNode<E> node = getNode(key);
        return node == null ? null : node.value;
    }

    private BSTNode<E> getNode(int key) {
        BSTNode<E> curr = root;
        while (curr != null) {
            if (key < curr.key) {
                curr = curr.left;
            } else if (key > curr.key) {
                curr = curr.right;
            } else {
                return curr;
            }
        }
        return null;
    }

    /**
     * 查找最小关键字对应值
     *
     * @return
     */
    public E min() {
        return min(root);
    }

    private E min(BSTNode<E> node) {
        BSTNode<E> curr = node;
        while (curr != null) {
            if (curr.left == null) {
                return curr.value;
            }
            curr = curr.left;
        }
        return null;

    }

    /**
     * 查找最大关键字对应值
     *
     * @return
     */
    public E max() {
        return max(root);
    }

    private E max(BSTNode<E> node) {
        BSTNode<E> curr = node;
        while (curr != null) {
            if (curr.right == null) {
                return curr.value;
            }
            curr = curr.right;
        }
        return null;
    }

    /**
     * 存储关键字和对应值
     * 简易的实现，不能保证自平衡
     * 如果 put 的 key 是依次递增，则树结构便是一根直线，跟数组类似了！
     *
     * @param key
     * @param value
     */
    public void put(int key, E value) {
        /*
            两种情况：key存在覆盖，不存在添加
         */
        root = putVal(key, value, root);
    }

    private BSTNode<E> putVal(int key, E value, BSTNode<E> node) {
        if (node == null) {
            return new BSTNode<>(key, value);
        }
        if (key < node.key) {
            node.left = putVal(key, value, node.left);
        } else if (key > node.key) {
            node.right = putVal(key, value, node.right);
        } else {
            node.value = value;
        }
        return node;
    }

    /**
     * 查找关键字的前驱值，如图：./img/img_1.png
     * 节点有左子树：当前key所在节点其左子树中key最大的节点值
     * 节点没有左子树：返回离它最近且自左而来的祖先节点
     *
     * @param key
     * @return 前驱节点值，如果返回 null 表示未找到或没有前驱节点
     */
    public E predecessor(int key) {
        BSTNode<E> curr = root;
        // 记录 key 最大且自左向右而来的祖先节点
        BSTNode<E> ancestorFormLeft = null;
        while (curr != null) {
            if (key < curr.key) {
                curr = curr.left;
            } else if (key > curr.key) {
                ancestorFormLeft = curr;
                curr = curr.right;
            } else {
                break;
            }
        }
        if (curr == null) {
            // 未找到
            return null;
        }

        // 找到了
        if (curr.left != null) {
            // 再寻找左侧最大值的节点
            return max(curr.left);
        }
        // 没有左子节点，那就找祖先节点
        // 如果祖先为 null，只有一种情况，那就是该 curr 节点是最靠左的节点
        return ancestorFormLeft == null ? null : ancestorFormLeft.value;
    }

    /**
     * 查找关键字的后继值
     *
     * @param key
     * @return
     */
    public E successor(int key) {
        BSTNode<E> curr = root;
        // 记录 key 最大且自右向左而来的祖先节点
        BSTNode<E> ancestorFormRight = null;
        while (curr != null) {
            if (key < curr.key) {
                ancestorFormRight = curr;
                curr = curr.left;
            } else if (key > curr.key) {
                curr = curr.right;
            } else {
                break;
            }
        }
        if (curr == null) {
            return null;
        }
        if (curr.right != null) {
            return min(curr.right);
        }
        return ancestorFormRight == null ? null : ancestorFormRight.value;
    }

    /**
     * 根据关键字删除
     * 参考图片：img\deleteNode.png
     *
     * @param key
     * @return
     */
    public E delete(int key) {
        BSTNode<E> del = root;
        // 记录 del 的父节点
        BSTNode<E> parent = null;
        while (del != null) {
            if (key < del.key) {
                parent = del;
                del = del.left;
            } else if (key > del.key) {
                parent = del;
                del = del.right;
            } else {
                break;
            }
        }
        if (del == null) {
            return null;
        }

        /*
            4种情况：
                - 删除节点没有左孩子，将右孩子托孤
                - 删除节点没有右孩子，将左孩子托孤
                - 删除节点左右孩子都没有，上面其中一种情况就可以处理
                - 删除节点左右孩子都有
         */

        if (del.left == null) {
            shift(parent, del, del.right);
        } else if (del.right == null) {
            shift(parent, del, del.left);
        } else {
            // 左右孩子都有的情况

            /*
                步骤：
                - 找出被删除节点的后继节点
                - 移出后继节点
                - 将后继节点代替被删除的节点
            */
            BSTNode<E> s = del.right;
            // 后继节点的父节点
            BSTNode<E> sParent = del;
            while (s.left != null) {
                sParent = s;
                s = s.left;
            }
            // s 就是后继节点

            // 判断后继节点是否相邻
            if (sParent != del) {
                // 不相邻，移出后继节点，并将删除节点的右边赋值给后继节点
                shift(sParent, s, s.right); // 能称为后继节点，那么他就不可能有左子节点
                s.right = del.right;
            }
            // 移除 del 节点，后继节点代替
            shift(parent, del, s);
            // 将后继节点替代掉删除的节点
            s.left = del.left; // s 就替代了被删除的 del 节点
        }

        return del.value;
    }

    /**
     * 托孤
     *
     * @param parent  待删除节点的父节点
     * @param deleted 待删除节点
     * @param child   被顶上去的节点
     */
    private void shift(BSTNode<E> parent, BSTNode<E> deleted, BSTNode<E> child) {
        if (parent == null) {
            root = child;
        } else if (parent.left == deleted) {
            parent.left = child;
        } else if (parent.right == deleted) {
            parent.right = child;
        }
    }

    /**
     * 根据关键字删除（递归实现）
     *
     * @param key
     * @return
     */
    public E delete2(int key) {
        ArrayList<E> result = new ArrayList<>();
        root = doDelete(key, root, result);
        return result.isEmpty() ? null : result.get(0);
    }

    /**
     * 递归删除
     *
     * @param key
     * @param node
     * @param result 保存被删除的元素，供后面处理
     * @return
     */
    private BSTNode<E> doDelete(int key, BSTNode<E> node, ArrayList<E> result) {
        if (node == null) {
            return null;
        }
        if (key < node.key) {
            // 表示key小于当前节点，继续向左找，返回值表示它的左节点，如果已经找到，那么返回值就是被删除后顶替的节点
            node.left = doDelete(key, node.left, result);
            return node;
        }
        if (key > node.key) {
            // 表示key小于当前节点，继续向左找
            node.right = doDelete(key, node.right, result);
            return node;
        }

        // 找到了
        result.add(node.value);
        if (node.left == null) {
            // 没有左节点，直接托孤右节点
            return node.right;
        }
        if (node.right == null) {
            return node.left;
        }

        // 查找后继节点
        BSTNode<E> s = node.right;
        while (s.left != null) {
            s = s.left;
        }
        // 再次递归
        s.right = doDelete(s.key, node.right, result);
        s.left = node.left;
        // 接手被删除节点的左边
        return s;
    }

    /**
     * 返回所有比 k 小的值 less
     *
     * @param k
     * @return
     */
    public List<E> lt(int k) {
        return less(k, false);
    }

    /**
     * @param itself 是否包含本身
     * @return
     */
    private List<E> less(int k, boolean itself) {
        List<E> list = new ArrayList<>();
        BSTNode<E> curr = root;
        // 查找到的元素值 等于 null 表示没有找到该节点
        LinkedList<BSTNode<E>> stack = new LinkedList<>();
        while (curr != null || !stack.isEmpty()) {
            if (curr != null) {
                stack.push(curr);
                curr = curr.left;
            } else {
                BSTNode<E> pop = stack.pop();
                if (pop.key < k) {
                    list.add(pop.value);
                } else if (k == pop.key && itself) {
                    list.add(pop.value);
                    break;
                } else {
                    break;
                }
                curr = pop.right;
            }
        }
        return list;
    }

    /**
     * 返回所有比 k 小或等于 k 的值
     *
     * @param k
     * @return
     */
    public List<E> le(int k) {
        return less(k, true);
    }

    /**
     * 返回所有比 k 大的值 greater
     *
     * @param k
     * @return
     */
    public List<E> gt(int k) {

        // 中序遍历实现
//        List<E> list = new ArrayList<>();
//        BSTNode<E> curr = root;
//        LinkedList<BSTNode<E>> stack = new LinkedList<>();
//        while (curr != null || !stack.isEmpty()) {
//            if (curr != null) {
//                stack.push(curr);
//                curr = curr.left;
//            } else {
//                BSTNode<E> pop = stack.pop();
//                if (pop.key > k) {
//                    list.add(pop.value);
//                }
//                curr = pop.right;
//            }
//        }

        // 反向中序遍历实现
        LinkedList<E> list = new LinkedList<>();
        BSTNode<E> curr = root;
        LinkedList<BSTNode<E>> stack = new LinkedList<>();
        while (curr != null || !stack.isEmpty()) {
            if (curr != null) {
                stack.push(curr);
                curr = curr.right;
            } else {
                BSTNode<E> pop = stack.pop();
                if (pop.key > k) {
                    list.addFirst(pop.value);
                } else {
                    break;
                }
                curr = pop.left;
            }
        }
        return list;
    }

    /**
     * 范围查询    k1 <= 值  <= k2
     *
     * @param
     * @return
     */
    public List<E> between(int k1, int k2) {
        if (k2 < k1) {
            throw new IllegalArgumentException("k2 must be greater than k1");
        }

        List<E> list = new ArrayList<>();
        BSTNode<E> curr = root;
        LinkedList<BSTNode<E>> stack = new LinkedList<>();
        while (curr != null || !stack.isEmpty()) {
            if (curr != null) {
                stack.push(curr);
                curr = curr.left;
            } else {
                BSTNode<E> pop = stack.pop();
                if (k1 <= pop.key && k2 >= pop.key) {
                    list.add(pop.value);
                } else if (pop.key > k2) {
                    // 提前结束
                    break;
                }
                curr = pop.right;
            }
        }

        return list;
    }

    /**
     * 验证 tree 是否满足二叉搜索树的特征
     * 方式一 中序遍历实现
     */
    public static <T> boolean isValidBST(BSTree<T> tree) {
        BSTNode<T> curr = tree.root;
        // 保存父节点的 key
        long prev = Long.MIN_VALUE;
        LinkedList<BSTNode<T>> stack = new LinkedList<>();
        while (curr != null || !stack.isEmpty()) {
            if (curr != null) {
                stack.push(curr);
                curr = curr.left;
            } else {
                BSTNode<T> pop = stack.pop();
                if (prev >= pop.key) {
                    return false;
                }
                prev = pop.key;
                curr = pop.right;
            }
        }
        return true;
    }

    /**
     * 方式二 递归实现
     */
    public static <T> boolean isValidBST2(BSTree<T> tree) {
        return validBST(tree.root, Long.MIN_VALUE, Long.MAX_VALUE);
    }

    private static <T> boolean validBST(BSTNode<T> node, long min, long max) {
        if (node == null) {
            return true;
        }
        if (node.key <= min || node.key >= max) {
            return false;
        }
        return validBST(node.left, min, node.key) && validBST(node.right, node.key, max);
    }

    /**
     * 树是否为空
     *
     * @return
     */
    public boolean isEmpty() {
        return root == null;
    }

    /**
     * 树节点
     *
     * @param <E>
     */
    public static class BSTNode<E> {
        public int key;
        public E value;
        public BSTNode<E> left;
        public BSTNode<E> right;

        public BSTNode(int key) {
            this.key = key;
        }

        public BSTNode(int key, E value) {
            this.key = key;
            this.value = value;
        }

        public BSTNode(E val, int key, BSTNode<E> left, BSTNode<E> right) {
            this.value = val;
            this.key = key;
            this.left = left;
            this.right = right;
        }

        @Override
        public String toString() {
            return "{" + key + ", " + value + "}";
        }
    }


    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        printVerticalTree(root, "", sb);
        return sb.toString();
    }

    private void printVerticalTree(BSTNode<E> node, String prefix, StringBuilder sb) {
        if (node == null) {
            return;
        }

        printVerticalTree(node.right, prefix + "        ", sb);
        sb.append(prefix).append("|--").append(node).append("\n");
        printVerticalTree(node.left, prefix + "        ", sb);
    }
}
